/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <sys/select.h>
#include <linux/vm_sockets.h>
#include <glib.h>
#include <sys/epoll.h>
#include <glib/gprintf.h>
#include <sys/syscall.h>
#include <netdb.h>
#include <stdbool.h>

#include "transfer.h"

#define EPOLL_SIZE          4
#define THREAD_SIZE         64
#define EPOLL_WAIT_BLOCK    (-1)    /* epoll event wait time，unit is millisecond, value[-1] means block */
#define SOCKET_WAIT_TIME    1       /* max time to wait for socket recv or send, unit is second */
#define MIN_LOCAL_PORT      1024
#define HOST_NAME_SIZE      64
#define REMOTE_PORT_LEN     128
#define MAX_PORT_VALUE      65535

gboolean is_running;
GThreadPool *gpool;

static gint local_port;
static gchar* remote_addr;
static gint remote_port;
static gboolean ipv6_config;

static GOptionEntry entries[] =
{
    { "local-port", 'l', 0, G_OPTION_ARG_INT, &local_port, "port of local listen", NULL },
    { "remote-addr", 'a', 0, G_OPTION_ARG_STRING, &remote_addr, "address of remote server", NULL },
    { "remote-port", 'p', 0, G_OPTION_ARG_INT, &remote_port, "port of remote server", NULL },
    { "ipv6", '6', 0, G_OPTION_ARG_NONE, &ipv6_config, "use only ipv6 address, optional", NULL },
    { NULL }
};

static int set_socket_wait_time(int fd)
{
    struct timeval timeout;
    int ret;

    timeout.tv_sec = SOCKET_WAIT_TIME;
    timeout.tv_usec = 0;
    ret = setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, (char *)&timeout, sizeof(struct timeval));
    if (ret < 0) {
        log_err("set fd = %d send timeout fail. message: %m", fd);
        goto err;
    }
    ret = setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, (char *)&timeout, sizeof(struct timeval));
    if (ret < 0) {
        log_err("set fd = %d recv timeout fail. message: %m", fd);
        goto err;
    }

err:
    return ret;
}

static bool check_event(uint32_t event)
{
    if ((event & EPOLLERR) || (event & EPOLLHUP) || (!(event & EPOLLIN))) {
        return -1;
    }
    return 0;
}

static bool transfer_result_func(int ret, int epoll_fd, int listen_fd,
                                 struct proxy_data *arg, int *server_pfd)
{
    bool epoll_flag = TRUE;
    switch (ret) {
        case TRAN_FAIL: {
            log_err("transfer data fail. close both client and server fd.");
            epoll_flag = FALSE;
            break;
        }
        case TRAN_CLOSE: {
            log_info("remote connection of source fd is closed. close both client and server fd.");
            epoll_flag = FALSE;
            break;
        }
        case TRAN_COMPLETE: {
            log_info("transfer message is complete.");
            if (listen_fd == arg->client_fd) {
                /* client is closed. */
                log_info("client fd is closed. close both client and server fd.");
                epoll_flag = FALSE;
                break;
            }
            log_info("server request is complete. keep server fd connecting."
                     " server fd = %d, client fd = %d", *server_pfd, arg->client_fd);
            break;
        }
        default: {
            /* transfer successful. do nothing. */
            break;
        }
    }

    return epoll_flag;
}

static bool epoll_event_func(int epoll_fd, int nfds, struct epoll_event *events_q,
                             struct proxy_data *arg, int *server_pfd)
{
    int i, ret;
    bool epoll_flag = TRUE;
    int client_fd = arg->client_fd;

    /*
    * There are only two handles(client_fd and server_fd) managed by this epoll_fd,
    * if errors occured on any of them, just close both and exit from this thread
    */
    for (i = 0; i < nfds; i++) {
        int listen_fd = events_q[i].data.fd;
        if (check_event(events_q[i].events)) {
            /* An error has occured on this fd, or this socket is not ready for reading */
            log_err("epoll event error. fd = %d. events = %d. %m", listen_fd, events_q[i].events);
            epoll_flag = FALSE;
            break;
        } else {
            log_info("transfer_data begin: i = %d, listen_fd = %d, local_port = %d",
                     i, listen_fd, local_port);
            /* Data on the fd waiting to be read. */
            if (listen_fd == client_fd) {
                ret = transfer_data(client_fd, *server_pfd);
            } else {
                ret = transfer_data(*server_pfd, client_fd);
            }
            log_info("transfer_data end: ret = %d, i = %d, listen_fd = %d, local_port = %d",
                     ret, i, listen_fd, local_port);
            epoll_flag = transfer_result_func(ret, epoll_fd, listen_fd, arg, server_pfd);
        }
    }

    return epoll_flag;
}

static void proxy_running(int epoll_fd, struct epoll_event *events_q,
                          struct proxy_data *arg, int *server_fd)
{
    gboolean listen_flag = TRUE;
    int nfds;

    while (listen_flag) {
        log_info("epoll_wait begin");

        nfds = epoll_wait(epoll_fd, events_q, EPOLL_SIZE, EPOLL_WAIT_BLOCK);
        if (nfds < 0) {
            log_err("epoll wait error. %m");
            continue;
        }
        listen_flag = epoll_event_func(epoll_fd, nfds, events_q, arg, server_fd);

        log_info("epoll_wait end");
    }
}

static int epoll_event_prepare(struct proxy_data *arg, int epoll_fd, int server_fd)
{
    int ret;
    struct epoll_event ev;

    log_info("add client fd in epoll. client fd = %d, local_port = %d, remote addr = %s",
             arg->client_fd, local_port, remote_addr);

    /* Register epoll event of client */
    ev.data.fd = arg->client_fd;
    ev.events = EPOLLIN;
    ret = epoll_ctl(epoll_fd, EPOLL_CTL_ADD, arg->client_fd, &ev);
    if (ret < 0) {
        log_err("epoll_ctl add client fd error! %m");
        goto out;
    }

    log_info("add server fd in epoll. server fd = %d, local_port = %d, remote addr = %s",
             server_fd, local_port, remote_addr);

    /* Register epoll event of server */
    ev.data.fd = server_fd;
    ev.events = EPOLLIN | EPOLLET;
    ret = epoll_ctl(epoll_fd, EPOLL_CTL_ADD, server_fd, &ev);
    if (ret < 0) {
        log_err("epoll_ctl add server fd error! %m");
        goto out;
    }

    /* Set wait time of send and recv for both client and server */
    ret = set_socket_wait_time(arg->client_fd);
    if (ret < 0) {
        log_err("set client socket wait time error! %m");
        goto out;
    }

    ret = set_socket_wait_time(server_fd);
    if (ret < 0) {
        log_err("set server socket wait time error! %m");
        goto out;
    }

out:
    return ret;
}

static gint proxy_func(gpointer data)
{
    struct proxy_data *arg = (struct proxy_data *)data;

    int epoll_fd, ret, server_fd;
    struct epoll_event *events_q;

    log_info("proxy_func begin. local_port = %d, remote addr = %s",
             local_port, remote_addr);

    epoll_fd = epoll_create(EPOLL_SIZE);
    if (epoll_fd < 0) {
        log_err("epoll_create fail.");
        goto free;
    }
    events_q = g_new0(struct epoll_event, EPOLL_SIZE);

    /* Create connection to server */
    server_fd = create_server_connect(arg);
    if (server_fd < 0) {
        log_err("create_server_connect error");
        goto free_events;
    }

    ret = epoll_event_prepare(arg, epoll_fd, server_fd);
    if (ret < 0) {
        log_err("epoll event prepare fail");
        goto free_server;
    }

    log_info("proxy_running begin. local_port = %d, remote addr = %s",
             local_port, remote_addr);

    proxy_running(epoll_fd, events_q, arg, &server_fd);

    log_info("listen thread is exiting.");
free_server:
    log_info("close server fd = %d", server_fd);
    close(server_fd);
free_events:
    g_free(events_q);
    log_info("close epooll fd = %d", epoll_fd);
    close(epoll_fd);
free:
    log_info("close client fd = %d", arg->client_fd);
    close(arg->client_fd);
    g_free(arg);
    return 0;
}

static void accept_connection(int local_fd, struct listen_data *arg)
{
    struct sockaddr_vm client_addr;
    socklen_t socklen = sizeof(struct sockaddr_vm);
    GError *err = NULL;
    int client_fd, ret = 0;

    memset(&client_addr, 0, sizeof(struct sockaddr_vm));

    while (1) {
        client_fd = accept(local_fd, (struct sockaddr *)&client_addr, &socklen);
        if (client_fd < 0) {
            if ((errno == EAGAIN) || (errno == EWOULDBLOCK)) {
                /* Processed all incoming connections */
                log_info("accept complete, local_port = %d: %m", local_port);
                break;
            }
            log_err("accept error, local_port = %d: %m", local_port);
            break;
        }
        log_info("accept success. local_port = %d, client_fd = %d, remote addr = %s",
                 local_port, client_fd, remote_addr);

        /* Init data for proxy, set remote addr and port */
        struct proxy_data *pri_data = (struct proxy_data *)g_malloc0(sizeof(struct proxy_data));
        pri_data->client_fd = client_fd;
        pri_data->remote_addr = arg->remote_addr;
        pri_data->rp = arg->rp;

        struct thread_func_data *thread_data =
                (struct thread_func_data *)g_malloc0(sizeof(struct thread_func_data));
        thread_data->data = pri_data;
        thread_data->func = proxy_func;
        /* Execute the transfer between client and server in a single thread */
        ret = g_thread_pool_push(gpool, (gpointer)(thread_data), &err);
        log_info("thread pool push result: ret = %d, err = %p", ret, err);
        if (err != NULL) {
            log_err("thread pool push error: %s", err->message);
            g_clear_error(&err);
            continue;
        }
    }
}

static void listen_func(struct listen_data *arg)
{
    log_info("local_fd = %d", arg->local_fd);

    int nfds;
    int epoll_fd = arg->epoll_fd;
    int local_fd = arg->local_fd;
    struct epoll_event *events_q = arg->events_q;

    while (is_running) {
        nfds = epoll_wait(epoll_fd, events_q, EPOLL_SIZE, EPOLL_WAIT_BLOCK);
        if (nfds < 0) {
            log_err("epoll wait error. %m");
            continue;
        }
        int i;
        for (i = 0; i < nfds; i++) {
            int listen_fd = events_q[i].data.fd;
            if (check_event(events_q[i].events)) {
                /* An error has occured on this fd, or this socket is not ready for reading */
                log_err("epoll event error. fd = %d . %m", listen_fd);
                is_running = FALSE;
                continue;
            } else {
                /* A notification on the listening socket, which means one or more incoming connections */
                accept_connection(local_fd, arg);
            }
        }
    }

    return;
}

static void thread_execute(gpointer data, gpointer user_data)
{
    log_info("thread execute %ld in", syscall(__NR_gettid));
    struct thread_func_data *arg = (struct thread_func_data *)data;
    (*(arg->func))(arg->data);
    g_free(data);
    log_info("thread execute %ld out", syscall(__NR_gettid));
}

static int cmd_line_parse(int argc, char *argv[], GOptionContext *context)
{
    GError *error = NULL;

    g_option_context_add_main_entries (context, entries, NULL);
    g_option_context_set_summary(context,
        "command example: \r\n./qt_proxy -l 8000 -a kms.br-iaas-odin1.myhuaweicloud.com -p 443");

    if  (!g_option_context_parse(context, &argc, &argv, &error)) {
        log_err("option parsing failed: %s", error->message);
        g_clear_error(&error);
        return -1;
    }

    if (remote_addr == NULL) {
        log_err("invalid param: remote_addr is NULL");
        return -1;
    }

    if (local_port <= MIN_LOCAL_PORT || local_port > MAX_PORT_VALUE ||
        remote_port <= 0 || remote_port > MAX_PORT_VALUE) {
        log_err("invalid param: local_port = %d, remote_addr = %s, remote_port = %d",
            local_port, remote_addr, remote_port);
        return -1;
    }

    return 0;
}

static int get_ip_address_info(struct addrinfo **addr_result, char *ip_addr_buf, int buff_len)
{
    int ret;
    struct addrinfo hints;
    struct addrinfo *result, *rp;
    char remote_port_str[REMOTE_PORT_LEN];

    memset(&hints, 0, sizeof(hints));
    memset(remote_port_str, 0, REMOTE_PORT_LEN);

    if (snprintf(remote_port_str, REMOTE_PORT_LEN, "%d", remote_port) < 0)
        return -1;

    log_info("ipv6_config = %d, remote_port_str = %s", ipv6_config, remote_port_str);
    if (ipv6_config) {
        hints.ai_family = AF_INET6;
    } else {
        hints.ai_family = AF_INET;
    }
    hints.ai_socktype = 0;
    hints.ai_flags = 0;
    hints.ai_protocol = 0;

    ret = getaddrinfo(remote_addr, remote_port_str, &hints, &result);
    if (ret) {
        log_err("can not get ip by hostname: %s, port: %s", remote_addr, remote_port_str);
        return -1;
    }

    for (rp = result; rp != NULL; rp = rp->ai_next) {
        switch (rp->ai_family) {
            case AF_INET: {
                struct sockaddr_in *ipv4;
                ipv4 = (struct sockaddr_in *)rp->ai_addr;
                inet_ntop(rp->ai_family, &ipv4->sin_addr, ip_addr_buf, buff_len);
                break;
            }
            case AF_INET6: {
                struct sockaddr_in6 *ipv6;
                ipv6 = (struct sockaddr_in6 *)rp->ai_addr;
                inet_ntop(rp->ai_family, &ipv6->sin6_addr, ip_addr_buf, buff_len);
                break;
            }
            default: {
                log_err("invalid socket type : ai_family = %d", rp->ai_family);
                freeaddrinfo(result);
                return -1;
            }
        }

        log_info("[IPv%d]HostName : %s, IP Address : %s, port = %d", rp->ai_family == AF_INET ? 4 : 6,
                 remote_addr, ip_addr_buf, remote_port);
        break;
    }

    *addr_result = result;

    return ret;
}

static int proxy_server_set(struct listen_data *pri_data_ptr, int epoll_fd, int local_fd,
                            struct epoll_event *events_q)
{
    gpool = g_thread_pool_new((GFunc)thread_execute, NULL, THREAD_SIZE, FALSE, NULL);
    if (gpool == NULL) {
        log_err("Can not create thread!\n");
        return -1;
    }
    pri_data_ptr->epoll_fd = epoll_fd;
    pri_data_ptr->local_fd = local_fd;
    pri_data_ptr->events_q = events_q;
    return 0;
}

static int proxy_server_start_listen(struct sockaddr_vm *addr, int local_fd)
{
    memset(addr, 0, sizeof(struct sockaddr_vm));
    addr->svm_family = AF_VSOCK;
    addr->svm_port = local_port;
    addr->svm_cid = VMADDR_CID_ANY;

    if (bind(local_fd, (struct sockaddr *)addr, sizeof(struct sockaddr_vm)) == -1) {
        log_err("bind error: %m");
        return -1;
    }

    if (listen(local_fd, THREAD_SIZE) < 0) {
        log_err("listen error: %m");
        return -1;
    }
    log_info("listen success");
    return 0;
}

static int proxy_server_prepare(struct listen_data *pri_data_ptr, struct sockaddr_vm *addr)
{
    int epoll_fd, local_fd;
    int ret;
    struct epoll_event *events_q = NULL;
    struct epoll_event ev;

    /* Init epoll event */
    epoll_fd = epoll_create(EPOLL_SIZE);
    if (epoll_fd < 0) {
        log_err("epoll_create fail.");
        return -1;
    }
    events_q = g_new0(struct epoll_event, EPOLL_SIZE);

    /* Create listen socket, set nonblock */
    local_fd = socket(AF_VSOCK, SOCK_STREAM | SOCK_NONBLOCK, 0);
    if (local_fd < 0) {
        log_err("create local socket fail.");
        goto free_queue;
    }

    ev.data.fd = local_fd;
    ev.events = EPOLLIN | EPOLLET;
    /* Register epoll event of local listening */
    ret = epoll_ctl(epoll_fd, EPOLL_CTL_ADD, local_fd, &ev);
    if (ret < 0) {
        log_err("epoll_ctl error!");
        goto free_local;
    }

    /* Init socket and start listen */
    if (proxy_server_start_listen(addr, local_fd) < 0) {
        goto free_local;
    }

    /* Create thread pool and init listen thread */
    if (proxy_server_set(pri_data_ptr, epoll_fd, local_fd, events_q) < 0) {
        goto free_local;
    }

    return 0;

free_local:
    close(local_fd);
free_queue:
    g_free(events_q);
    close(epoll_fd);

    return -1;
}

static void proxy_server_release(struct listen_data *pri_data_ptr)
{
    g_thread_pool_free(gpool, FALSE, TRUE);
    gpool = NULL;
    close(pri_data_ptr->local_fd);
    g_free(pri_data_ptr->events_q);
    close(pri_data_ptr->epoll_fd);
}

int main(int argc, char *argv[])
{
    struct sockaddr_vm addr;
    struct listen_data pri_data;
    struct addrinfo *result;
    char ip_addr_buf[HOST_NAME_SIZE];
    int ret;

    /* command line option parse */
    GOptionContext *context;

    context = g_option_context_new ("- qt proxy tree model performance");
    ret = cmd_line_parse(argc, argv, context);
    if (ret) {
        goto free;
    }

    /* get ip address from hostname */
    ret = get_ip_address_info(&result, ip_addr_buf, HOST_NAME_SIZE);
    if (ret) {
        log_err("get_ip_address_info fail");
        goto free;
    }

    /* create epoll event and init socket */
    ret = proxy_server_prepare(&pri_data, &addr);
    if (ret) {
        log_err("proxy_server_prepare fail");
        goto free_addr;
    }

    pri_data.remote_addr = ip_addr_buf;
    pri_data.rp = result;

    is_running = TRUE;
    listen_func(&pri_data);

    /* Wait until all threads are complete */
    proxy_server_release(&pri_data);

free_addr:
    freeaddrinfo(result);
free:
    g_option_context_free(context);
    return 0;
}
